FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention
FlashMask-3：高效掩码感知分布式注意力
🧭 训练系统 · 掩码感知注意力 ★ 深度解读
Baidu Inc.
将列式稀疏掩码扩展至 Hopper 分布式注意力，实现零通信掩码本地化与负载均衡的长上下文训练加速
📌 工作参考价值（4 点）
- 长序列行为建模或多文档检索中，采用列式稀疏掩码（column-wise sparse mask）编码 document/prefix-LM/sliding-window 等结构化 attention，用 O(N) 边界向量替代 O(N^2) 显式掩码，kernel 跳过全掩码 tile，可大幅降低显存和计算，适合用户长期行为序列的 Transformer 训练。
- 分布式 context parallel 下，掩码边界向量仅需做 clip–shift–max 变换即得到 rank-local 掩码，零跨卡通信；在电商/搜索推荐的多机训练中可直接复用，避免 mask 同步开销。
- 对稀疏 mask 导致的负载不均，借鉴 Hierarchical Compute-Communication IPO：先 IPO 做 compute balance，再 inter-node 交换降低 KV 通信量，最后 intra-node rebalance，在工程上可分层优化。
- 通信重叠可使用 NVSHMEM 对称内存 + 一端 RDMA + mask 驱动的 chunk 剪枝，减少约 50% all-gather/reduce-scatter 流量，提升 overlap 效率，适用于大规模分布式训练（如 LLM 驱动的搜索推荐）。